#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UCEIS评分模型 - 兼容现有inference_server.py的接口
"""

import torch
import torch.nn as nn
import torchvision.models as models
from typing import Dict, Any


class UCEISScoreModel(nn.Module):
    """
    UCEIS评分模型 - 兼容现有推理服务器
    模拟多分类输出格式
    """

    def __init__(self, config: Dict[str, Any]):
        super().__init__()

        self.config = config
        self.num_classes = config.get('num_classes', 8)
        self.feature_dims = config.get('feature_dims', [64, 128, 256, 512])

        # 检查是否启用多模态
        self.use_multimodal = config.get('use_multimodal', False)

        if self.use_multimodal:
            # 多模态版本 - 从训练好的图像分支模型加载
            self.setup_multimodal_model()
        else:
            # 单模态版本 - 基于ResNet
            self.setup_resnet_model()

    def setup_resnet_model(self):
        """设置基于ResNet的模型"""
        # 使用ResNet50作为backbone
        self.backbone = models.resnet50(pretrained=True)

        # 替换最后的分类层
        num_ftrs = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Linear(num_ftrs, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, self.num_classes)
        )

    def setup_multimodal_model(self):
        """设置多模态模型（从训练好的权重加载）"""
        # 这里应该从训练好的ImageOnlyUCModel加载权重
        # 为了兼容性，创建类似的架构

        self.backbone = models.resnet50(pretrained=False)

        # 替换最后的分类层
        num_ftrs = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(128, self.num_classes)
        )

    def forward(self, x):
        """
        前向传播

        Args:
            x: 输入图像张量 (B, 3, H, W)

        Returns:
            logits: 分类logits (B, num_classes)
        """
        logits = self.backbone(x)

        # 为了与现有inference_server.py兼容，返回字典格式
        return {
            'uceis_logits': logits
        }

    def load_from_image_only_checkpoint(self, checkpoint_path: str):
        """从仅图像模型加载权重"""
        print(f"从 {checkpoint_path} 加载权重...")

        try:
            checkpoint = torch.load(checkpoint_path, map_location='cpu')

            if 'model_state_dict' in checkpoint:
                state_dict = checkpoint['model_state_dict']
            else:
                state_dict = checkpoint

            # 尝试匹配权重
            model_dict = self.state_dict()

            # 过滤不匹配的键
            filtered_dict = {k: v for k, v in state_dict.items()
                           if k in model_dict and v.shape == model_dict[k].shape}

            # 加载匹配的权重
            model_dict.update(filtered_dict)
            self.load_state_dict(model_dict)

            print(f"成功加载 {len(filtered_dict)} 个权重参数")

        except Exception as e:
            print(f"加载权重失败: {e}")
            print("使用预训练权重作为后备")

    def predict_uceis_score(self, x):
        """预测UCEIS分数"""
        with torch.no_grad():
            logits = self.forward(x)
            if isinstance(logits, dict):
                logits = logits['uceis_logits']

            predicted_class = torch.argmax(logits, dim=1)
            # 转换为1-8分
            uceis_score = predicted_class + 1

            return uceis_score, logits


def create_uceis_model_from_checkpoint(checkpoint_path: str, config: Dict[str, Any] = None) -> UCEISScoreModel:
    """
    从检查点创建UCEIS模型

    Args:
        checkpoint_path: 模型检查点路径
        config: 模型配置

    Returns:
        配置好的UCEISScoreModel实例
    """
    if config is None:
        config = {
            'num_classes': 8,
            'feature_dims': [64, 128, 256, 512],
            'use_multimodal': True
        }

    model = UCEISScoreModel(config)

    # 尝试加载训练好的权重
    if checkpoint_path and checkpoint_path.endswith('.pth'):
        model.load_from_image_only_checkpoint(checkpoint_path)

    return model


def test_uceis_model():
    """测试UCEIS模型"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 创建模型配置
    config = {
        'num_classes': 8,
        'feature_dims': [64, 128, 256, 512],
        'use_multimodal': True
    }

    model = UCEISScoreModel(config).to(device)

    # 创建测试输入
    batch_size = 2
    test_input = torch.randn(batch_size, 3, 224, 224).to(device)

    print(f"测试输入尺寸: {test_input.shape}")

    # 前向传播
    with torch.no_grad():
        outputs = model(test_input)

        if isinstance(outputs, dict):
            logits = outputs['uceis_logits']
            print(f"输出字典格式: {list(outputs.keys())}")
        else:
            logits = outputs

        print(f"输出logits尺寸: {logits.shape}")

        # 预测分数
        scores, logits_out = model.predict_uceis_score(test_input)
        print(f"预测UCEIS分数: {scores}")

    print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")

    return model


if __name__ == "__main__":
    test_uceis_model()